"""
Case Type   : python驱动psycopg2
Case Name   : python驱动psycopg2公共函数
Create At   : 2023-08-23
Owner       : n@ningyali
Description :
    1、python驱动psycopg2公共函数
Expect      :
    1、python驱动psycopg2公共函数
History     :
"""
import os
import subprocess
import traceback

from yat.test import macro

from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger


class Psycopg2Common:
    def __init__(self, **kwargs):
        self.log = Logger()
        self.pri_sh = CommonSH('PrimaryDbUser')
        self.cons = Constant()
        self.pkg_path = os.path.dirname(macro.DB_SCRIPT_PATH)
        self.new_dir = "opg_psycopg2"
        self.cur_ip = subprocess.getoutput("hostname -I|awk '{{print $1}}'")
        self.dbname = kwargs['dbname']
        self.user = kwargs['user']
        self.password = kwargs['password']
        self.host = kwargs['host']
        self.port = kwargs['port']
        self.py_path = None

    def decompress_pkg(self):
        """
        解压***-Python.tar.gz得到lib和psycopg2
        :return: True执行成功，False执行失败
        """
        text = '----解压***-Python.tar.gz得到lib----'
        self.log.info(text)
        arch = subprocess.getoutput('arch')
        cmd = f"cd {self.pkg_path} && " \
            f"if [ -e opg_psycopg2 ]; then echo 'exists'; " \
            f"else mkdir opg_psycopg2; fi && " \
            f"if [ -e openGauss-Python-*-{arch}.tar.gz ]; " \
            f"then tar -zxf openGauss-Python-*-{arch}.tar.gz " \
            f"-C {self.new_dir} && " \
            f"ls -l {self.new_dir}; " \
            f"else echo 'PDBC package not exists'; " \
            f"fi "
        self.log.info(cmd)
        res = subprocess.getoutput(cmd)
        self.log.info(res)
        if 'lib' in res:
            return True
        else:
            return False

    def set_new_psycopg2(self):
        """
        使用root用户将psycopg2拷贝到python安装目录下的site-packages文件夹下，并修改psycopg2目录权限为755
        :return:True执行成功，False执行失败
        """
        flag = False
        cmd = "find /usr/lib -type d -name 'site-packages'|grep python3|" \
              "head -1"
        self.log.info(cmd)
        self.py_path = subprocess.getoutput(cmd)
        self.log.info(self.py_path)

        if self.py_path.strip():
            pdbc_path = os.path.join(self.pkg_path, self.new_dir, 'psycopg2')
            cmd = f"\\cp -r {pdbc_path} {self.py_path}/ && " \
                f"chmod 755 -R {os.path.join(self.py_path, 'psycopg2')}"
            self.log.info(cmd)
            res = subprocess.getoutput(cmd)
            self.log.info(res)
            flag = True
        return flag

    def set_lib_path(self):
        """
        将解压后的lib目录，配置在LD_LIBRARY_PATH中
        :return:True执行成功，False执行失败
        """
        text = '----配置lib,连接数据库----'
        self.log.info(text)
        lib_path = os.path.join(self.pkg_path, self.new_dir, "lib")
        os.environ['LD_LIBRARY_PATH'] += f':{lib_path}'
        os.environ['PYTHONPATH'] += f':{self.py_path}'
        cmd = f'env|grep LD_LIBRARY_PATH|grep "{lib_path}";' \
            f'echo $PYTHONPATH|grep "{self.py_path}"'
        self.log.info(cmd)
        res = subprocess.getoutput(cmd)
        self.log.info(res)
        if not res.strip():
            return False
        else:
            return True

    def bak_conf(self, bak_flag=True, f_name=None):
        """
        备份/还原pg_hba文件
        :param bak_flag: True备份，False还原
        :param f_name: 默认为pg_hba
        :return:True执行成功，False执行失败
        """
        is_right = False
        if not f_name:
            f_name = os.path.join(macro.DB_INSTANCE_PATH,
                                  macro.PG_HBA_FILE_NAME)
        else:
            f_name = os.path.join(macro.DB_INSTANCE_PATH, f_name)
        if bak_flag:
            cmd = f"\\cp {f_name} {f_name}_bak.bak"
        else:
            cmd = f"if [ -e {f_name}_bak.bak ]; " \
                f"then \\cp {f_name}_bak.bak {f_name}; " \
                f"rm -rf {f_name}_bak.bak; fi"
        self.log.info(cmd)
        res = self.pri_sh.node.sh(cmd).result()
        self.log.info(res)
        if not bak_flag and not res.strip():
            res1 = self.pri_sh.reload_db_config()
            if 'server signaled' in res1:
                is_right = True
        elif bak_flag and not res.strip():
            is_right = True
        return is_right

    def set_pghba(self, open_flag=True, db=None, user=None, mode='sha256'):
        """
        开启/取消白名单
        :param open_flag: True开启，False取消
        :param db: 数据库名
        :param user: 数据库用户名
        :param mode: 认证方式
        :return: True执行成功，False执行失败
        """
        pghba_path = os.path.join(macro.DB_INSTANCE_PATH,
                                  macro.PG_HBA_FILE_NAME)
        hostmode = 'host' if mode != 'cert' else 'hostssl'
        grep_cmd = f"grep {self.cur_ip} {pghba_path}|grep sha256|" \
            f"grep -v '#'|awk '{{{{print $2, $3}}}}'"
        param_list = list()
        set_mode = mode if open_flag else ''
        db_name = db if db else 'all'
        db_user = user if user else 'all'
        try:
            if mode == 'cert':
                pghba_param = f'{hostmode} {db_name} {db_user} ' \
                    f'{self.cur_ip}/32 {set_mode}'
                param_list.append(pghba_param)
                pghba_param = f'{hostmode} all all 127.0.0.1/32 {set_mode}'
                param_list.append(pghba_param)
            elif db or user:
                pghba_param = f'{hostmode} {db_name} {db_user} ' \
                    f'{self.cur_ip}/32 {set_mode}'
                param_list.append(pghba_param)
            else:
                grep_res = self.pri_sh.node.sh(grep_cmd).result()
                self.log.info(grep_res)
                if grep_res.strip():
                    for line in grep_res.strip().splitlines():
                        db_name = db if db else line.strip().split()[0]
                        db_user = user if user else line.strip().split()[1]
                        pghba_param = f'{hostmode} {db_name} {db_user} ' \
                            f'{self.cur_ip}/32 {set_mode}'
                        param_list.append(pghba_param)
            if param_list:
                self.log.info(param_list)
                str_params = ' '.join([f'-h "{i}"' for i in param_list])
                guc_cmd = f'source {macro.DB_ENV_PATH};' \
                    f'gs_guc reload -D {macro.DB_INSTANCE_PATH} ' \
                    f'{str_params}'
                self.log.info(guc_cmd)
                guc_res = self.pri_sh.node.sh(guc_cmd).result()
                self.log.info(guc_res)
                assert 'Failed instances: 0' in guc_res
            else:
                raise Exception('NotFound: 未找到记录，请检查pg_hba文件配置！')
        except Exception as e:
            grep_cmd = f"grep {self.cur_ip} {pghba_path}"
            grep_res = self.pri_sh.node.sh(grep_cmd).result()
            self.log.info(grep_res)
            self.log.error(f'str(e): {e}')
            self.log.error(f'repr(e): {repr(e)}')
            self.log.error(f'traceback.format_exc():\n'
                           f'{traceback.format_exc()}')
            return False
        else:
            grep_cmd = f"grep {self.cur_ip} {pghba_path}"
            grep_res = self.pri_sh.node.sh(grep_cmd).result()
            self.log.info(grep_res)
            return True

    @staticmethod
    def format_sql(sql, *args):
        """
        格式化sql语句，返回psycopg2可执行SQL语句
        :param sql: SQL语句
        :param args: SQL中传递参数序列
        :return: 返回psycopg2可执行SQL语句
        """
        if args:
            f_sql = f"cur.execute(\\\"{sql}\\\",{args})\n"
        else:
            f_sql = ''
            for s in sql.strip(';').split(';'):
                f_sql += f"cur.execute(\\\"{s}\\\")\n"
        return f_sql

    @staticmethod
    def format_sql_with_executemany(sql, *args):
        """
        executemany格式化sql语句，返回psycopg2可执行SQL语句
        :param sql: SQL语句
        :param args: SQL中传递参数序列
        :return: 返回psycopg2可执行SQL语句
        """
        f_sql = f"cur.executemany(\\\"{sql}\\\",{args})\n"
        return f_sql

    @staticmethod
    def join_run_sql(sql, rollback=None):
        """
        组合可执行的sql语句
        :param sql: sql语句
        :param rollback: None不回滚，implicit隐式回滚，explicit显式回滚
        :return:可执行的sql语句
        """
        f_sql = "cur=conn.cursor()\n"
        f_sql += f"{sql}\n"
        if not rollback:
            f_sql += "conn.commit()\n"
        elif rollback == 'explicit':
            f_sql += "conn.rollback()\n"
        if 'select' in sql.lower():
            f_sql += "results=cur.fetchall()\n"
            f_sql += "print(results)\n"
        f_sql += "cur.close()\n"
        return f_sql

    def run_sql(self, sql, by='default'):
        """
        执行sql语句，返回结果
        :param sql: sql语句
        :param by: 数据库连接方式
        :return: 连接实例
        """
        res = self.conn_db(by, sql)
        return res

    def conn_db(self, by='default', sql='', cert_path=None, ssl_mode=None):
        """
        连接数据库
        :param by:数据库连接方式
        :param sql:sql语句
        :param cert_path:by为ssl时，指定证书路径
        :param ssl_mode:ssl连接时sslmode值
        :return: 连接实例
        """
        if by == 'long_string':
            conn_info = f"\\\"dbname={self.dbname} " \
                f"user={self.user} " \
                f"password={self.password} " \
                f"host={self.host} " \
                f"port={self.port}\\\""
        elif by == 'ssl':
            conn_info = f"dbname=\\\"{self.dbname}\\\"," \
                f"user=\\\"{self.user}\\\"," \
                f"password=\\\"{self.password}\\\"," \
                f"host=\\\"{self.host}\\\"," \
                f"port=\\\"{self.port}\\\"," \
                f"sslmode=\\\"{ssl_mode}\\\"," \
                f"sslcert=\\\"{os.path.join(cert_path, 'client.crt')}\\\"," \
                f"sslkey=\\\"{os.path.join(cert_path, 'client.key')}\\\", " \
                f"sslrootcert=\\\"{os.path.join(cert_path, 'cacert.pem')}\\\""
        else:
            conn_info = f"dbname=\\\"{self.dbname}\\\"," \
                f"user=\\\"{self.user}\\\"," \
                f"password=\\\"{self.password}\\\"," \
                f"host=\\\"{self.host}\\\"," \
                f"port=\\\"{self.port}\\\""
        cmd = f"python3 -c \"import psycopg2\n" \
            f"conn = psycopg2.connect({conn_info})\n" \
            f"print(conn.dsn)\n" \
            f"{sql}\n" \
            f"conn.close()\""
        scp_cmd = f'''expect <<EOF
            set timeout 300
            spawn {cmd}
            expect {{
                "*Enter PEM pass phrase:*" \
                {{ send "{self.pri_sh.node.ssh_password}\\n" }}
                eof {{ send_user "exit\\n" }}
            }}\n''' + 'EOF'
        self.log.info(scp_cmd)
        res = subprocess.getoutput(scp_cmd)
        self.log.info(res)
        return res

    def clean_decompress_file(self):
        """
        清理解压缩文件
        :return:True执行成功，False执行失败
        """
        cmd = f"if [ -e {os.path.join(self.pkg_path, self.new_dir)} ]; " \
            f"then rm -rf {os.path.join(self.pkg_path, self.new_dir)}; " \
            f"fi "
        self.log.info(cmd)
        res = subprocess.getoutput(cmd)
        self.log.info(res)
        if not res.strip():
            return True
        else:
            return False
